import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import PowerNorm, LogNorm

from bm3d import bm3d
from scipy.ndimage import median_filter
from scipy.ndimage import convolve
from scipy import ndimage
from skimage.filters import sobel
from skimage.metrics import structural_similarity as ssim
from pyhdf.SD import SD, SDC

import os
import math
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from scipy.signal import windows
from tqdm import tqdm

def load_hdf_as_numpy(file_path, var_name):
    f = SD(file_path, SDC.READ)
    data = f.select(var_name).get()
    arr = np.array(data, dtype=np.float32)
    return arr

def print_results(data_filt, 
                  data_unf, 
                  lower_percentile=2,
                  upper_percentile=98,
                  cmap="inferno",
                  norm_type="linear",
                  gamma=0.5):
    
    # --- percentile-based vmin / vmax ---
    all_data = np.concatenate([data_filt.ravel(), data_unf.ravel()])
    diff = data_unf - data_filt
    vmin = np.percentile(all_data, lower_percentile)
    vmax = np.percentile(all_data, upper_percentile)

    # --- choose normalization ---
    if norm_type == "log":
        norm = LogNorm(vmin=max(vmin, 1e-6), vmax=vmax)
    elif norm_type == "power":
        norm = PowerNorm(gamma=gamma, vmin=vmin, vmax=vmax)
    else:  # linear
        norm = None

    # --- plot ---
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))

    im0 = axes[0].imshow(data_unf, cmap=cmap, vmin=vmin, 
                         vmax=vmax, norm=norm, aspect="auto")
    axes[0].set_title(f"UNFILTERED")

    im1 = axes[1].imshow(data_filt, cmap=cmap, vmin=vmin, 
                         vmax=vmax, norm=norm, aspect="auto")
    axes[1].set_title(f"FILTERED")

    dmin = np.percentile(diff, 1); dmax = np.percentile(diff, 99)
    im2 = axes[2].imshow(diff, cmap="seismic", vmin=dmin, vmax=dmax, 
                         norm=norm, aspect="auto")
    axes[2].set_title(f"DIFFERENCE (UNFILTERED - FILTERED)")

    fig.colorbar(im0, ax=axes, fraction=0.025)
    # plt.tight_layout()
    plt.show()

def evaluate_denoising(I: np.ndarray, D: np.ndarray, eps: float = 1e-10) -> dict:
    """
    I : original/noisy image (2D numpy array)
    D : denoised image (2D numpy array)

    Returns dictionary with:
    - Noise Reduction Rate
    - ENL increase
    - Edge Preservation Index (EPI)
    - High Frequency Energy Ratio
    - Structure in residual (SSIM)
    """

    I = I.astype(np.float64)
    D = D.astype(np.float64)

    residual = I - D

    # --- 1. Noise Reduction Rate
    var_I = np.var(I)
    var_R = np.var(residual)
    nrr = 1 - (var_R / (var_I + eps))

    # --- 2. Equivalent Number of Looks (ENL)
    def enl(im):
        m = np.mean(im)
        v = np.var(im)
        return (m ** 2) / (v + eps)

    enl_I = enl(I)
    enl_D = enl(D)
    enl_increase = (enl_D - enl_I) / (enl_I + eps)

    # --- 3. Edge Preservation Index (correlation of Sobel edges)
    G_I = sobel(I)
    G_D = sobel(D)

    epi_num = np.sum((G_I - G_I.mean()) * (G_D - G_D.mean()))
    epi_den = np.sqrt(np.sum((G_I - G_I.mean())**2) * np.sum((G_D - G_D.mean())**2)) + eps
    epi = epi_num / epi_den

    # --- 4. High Frequency Energy ratio
    def high_freq_energy(im):
        low = ndimage.gaussian_filter(im, sigma=2)
        high = im - low
        return np.sum(high ** 2)

    hf_I = high_freq_energy(I)
    hf_D = high_freq_energy(D)

    hf_ratio = hf_D / (hf_I + eps)

    # --- 5. Structural content in residual (SSIM)
    h, w = residual.shape
    ssim_residual = ssim(residual, np.zeros((h, w)), data_range=residual.max() - residual.min())

    return {
        "Noise reduction rate (ideal 0.4–0.85)": round(float(nrr), 4),
        "ENL original": round(float(enl_I), 3),
        "ENL denoised": round(float(enl_D), 3),
        "ENL relative increase (ideal > 1.0)": round(float(enl_increase), 4),
        "Edge preservation index (ideal 0.8–1.05)": round(float(epi), 4),
        "HF energy ratio (ideal 0.3–0.7)": round(float(hf_ratio), 4),
        "Residual structure SSIM (ideal ≈ 0)": round(float(ssim_residual), 6),
    }

filtered_file_name = 'metm24_TA_251128_0101_9074_01_01A.hdf'
unfiltered_file_name = 'metm24_TA_251128_0101_9074_01_02A.hdf'

filtered_var_name = "m_m24_13_53_80V"
unfiltered_var_name = "m_m2413_53_80V"

data_filt = load_hdf_as_numpy(filtered_file_name, filtered_var_name)
data_unf = load_hdf_as_numpy(unfiltered_file_name, unfiltered_var_name)

print_results(data_filt, data_unf)

# Save images
np.save('images\\initial_data_unfiltered.npy', data_unf)
np.save('images\\initial_data_filtered.npy', data_filt)

WINDOW_SIZE = 5   

data_unf = load_hdf_as_numpy(unfiltered_file_name, unfiltered_var_name)

print("Image shape:", data_unf.shape)
print("Data type:", data_unf.dtype)

Image shape: (1212, 140)
Data type: float32

filtered_image = median_filter(data_unf, size=WINDOW_SIZE)

print_results(filtered_image, data_unf)

# Save image
np.save('images\\baseline_median_data_filtered.npy', filtered_image)

data_unf = load_hdf_as_numpy(unfiltered_file_name, unfiltered_var_name)

# ======================
# NORMALIZE IMAGE TO [0,1]
# ======================

min_val = np.min(data_unf)
max_val = np.max(data_unf)

image_norm = (data_unf - min_val) / (max_val - min_val + 1e-8)

print("Normalized range:", image_norm.min(), "to", image_norm.max())

# ======================
# APPLY BM3D
# ======================

# =====================
# AUTO NOISE ESTIMATION
# =====================

def estimate_noise_sigma(image):
    """
    Robust noise estimation using median absolute deviation
    Works best on satellite images
    Returns sigma in normalized [0,1] space
    """
    H = np.array([[1, -2, 1],
                  [-2, 4, -2],
                  [1, -2, 1]])

    im = convolve(image, H, mode='nearest')
    sigma = np.median(np.abs(im)) / 0.6745
    return sigma


sigma_est = estimate_noise_sigma(image_norm)

print(f"Estimated noise sigma: {sigma_est:.5f}")


filtered_norm = bm3d(image_norm, sigma_psd=sigma_est)

# ======================
# DE-NORMALIZE BACK
# ======================

filtered_image = filtered_norm * (max_val - min_val) + min_val

Normalized range: 0.0 to 1.0
Estimated noise sigma: 0.17369

print_results(filtered_image, data_unf)

# Save images
np.save('images\\bm3d_data_filtered.npy', filtered_image)

"""
advanced_denoise_pipeline

Integrated pipeline: DnCNN, TinyNAFNet, Full NAFNet (paper-like) with GPU support.
- Loads HDF swath (your var), self-supervised finetune if no weights, patch-based batched inference on GPU,
- Mixed precision training/inference for speed & memory efficiency,
- Visualization & no-ground-truth evaluation measures.

Usage:
  - Place this file with your HDF file.
  - Optionally put pretrained weights in ./weights/dncnn.pth, ./weights/tinynaf.pth, ./weights/nafnet.pth
  - Run: python advanced_denoise_pipeline
"""

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

def run_experiment(model_name='dncnn'):
    # -------------------------
    # User / runtime parameters
    # -------------------------
    UNFILTERED_FILE = 'metm24_TA_251128_0101_9074_01_02A.hdf'
    UNFILTERED_VAR = "m_m2413_53_80V"
    
    WEIGHTS_DIR = "./weights"
    DN_CNN_WEIGHTS = os.path.join(WEIGHTS_DIR, "dncnn.pth")
    TINY_NAF_WEIGHTS = os.path.join(WEIGHTS_DIR, "tinynaf.pth")
    NAFNET_WEIGHTS = os.path.join(WEIGHTS_DIR, "nafnet.pth")
    
    # Choose model: 'dncnn', 'tinynaf', 'nafnet'
    MODEL_NAME = model_name
    
    # Patch & inference config (tunable)
    PATCH_H = 128
    PATCH_W = 128
    STRIDE_H = PATCH_H // 2
    STRIDE_W = PATCH_W // 2
    BATCH_INFERENCE = 1   # number of patches processed at once during inference (GPU batched)
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # DEVICE = torch.device('cpu')
    NUM_WORKERS = 0
    # PIN_MEMORY = True if DEVICE.type == "cuda" else False
    PIN_MEMORY= False
    torch.set_num_threads(4)
    # Finetune config (self-supervised) - small on GPU
    FINETUNE = True        # set False if you have supervised weights
    FINETUNE_EPOCHS = 60
    FINETUNE_BATCH = 16
    LEARNING_RATE = 1e-3
    
    # Mixed-precision
    USE_AMP = True if DEVICE.type == "cuda" else False
    # USE_AMP = False
    
    # -------------------------
    # Utilities
    # -------------------------
    # Можно это раскомментировать если надо считывать из hdf, 
    # но у меня не получилось сделать так, чтобы и pyhdf работал, и
    # pytorch с cuda GPU работал (что-то одно отваливается). Поэтому 
    # тут и далее внизу читаем и сохраняем в npy нормальный формат :)
    # def load_hdf_as_numpy(file_path, var_name):
    #     f = SD(file_path, SDC.READ)
    #     data = f.select(var_name).get()
    #     arr = np.array(data, dtype=np.float32)
    #     return arr
    
    def make_hanning_window(h, w):
        win_h = windows.hann(h, sym=False)
        win_w = windows.hann(w, sym=False)
        win2d = np.outer(win_h, win_w).astype(np.float32)
        win2d += 1e-6
        return win2d
    
    # -------------------------
    # Models
    # -------------------------
    # 1) DnCNN (small)
    class DnCNN(nn.Module):
        def __init__(self, in_channels=1, features=64, depth=12):
            super().__init__()
            layers = [nn.Conv2d(in_channels, features, 3, padding=1), nn.ReLU(inplace=True)]
            for _ in range(depth - 2):
                layers += [nn.Conv2d(features, features, 3, padding=1, bias=False),
                           nn.BatchNorm2d(features),
                           nn.ReLU(inplace=True)]
            layers += [nn.Conv2d(features, in_channels, 3, padding=1)]
            self.net = nn.Sequential(*layers)
        def forward(self, x):
            # noise = self.net(x)
            # return x - noise
            return self.net(x)
    
    # 2) Tiny NAF-style (light)
    class SimpleNAFBlock(nn.Module):
        def __init__(self, channels):
            super().__init__()
            self.conv1 = nn.Conv2d(channels, channels, 1)
            self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
            self.conv3 = nn.Conv2d(channels, channels, 1)
            self.act = nn.GELU()
            self.beta = nn.Parameter(torch.zeros(1, channels, 1, 1))
        def forward(self, x):
            y = self.conv1(x)
            y = self.conv2(y)
            y = self.act(y)
            y = self.conv3(y)
            return x + self.beta * y
    
    class TinyNAFNet(nn.Module):
        def __init__(self, in_ch=1, out_ch=1, width=32, num_blocks=6):
            super().__init__()
            self.entry = nn.Conv2d(in_ch, width, 3, padding=1)
            self.body = nn.Sequential(*[SimpleNAFBlock(width) for _ in range(num_blocks)])
            self.exit = nn.Conv2d(width, out_ch, 3, padding=1)
        def forward(self, x):
            fe = self.entry(x)
            fb = self.body(fe)
            out = self.exit(fb)
            # return x - out
            return out
    
    # 3) Full paper-like NAFNet (encoder-decoder U-Net style)
    # Implement core building blocks inspired by the NAFNet paper:
    # NAF building blocks (faithful but practical)
    class SimpleGate(nn.Module):
        def forward(self, x):
            a, b = x.chunk(2, dim=1)
            return a * b
    
    # Helper: LayerNorm2d (standard in NAFNet)
    class LayerNorm2d(nn.Module):
        def __init__(self, channels, eps=1e-6):
            super().__init__()
            self.weight = nn.Parameter(torch.ones(1, channels, 1, 1))
            self.bias = nn.Parameter(torch.zeros(1, channels, 1, 1))
            self.eps = eps
    
        def forward(self, x):
            mean = x.mean(dim=(1, 2, 3), keepdim=True)
            var = x.var(dim=(1, 2, 3), keepdim=True)
            x = (x - mean) / (var + self.eps).sqrt()
            return x * self.weight + self.bias
    
    class SimpleChannelAttention(nn.Module):
        def __init__(self, channels):
            super().__init__()
            self.avg_pool = nn.AdaptiveAvgPool2d(1)
            self.conv = nn.Conv2d(channels, channels, 1, bias=True)
        def forward(self, x):
            v = self.avg_pool(x)
            v = self.conv(v)
            return x * torch.sigmoid(v)
    
    class NAFBlockV2(nn.Module):
        def __init__(self, channels):
            super().__init__()
            # First PW -> DW -> SimpleGate -> PW
            self.pw1 = nn.Conv2d(channels, channels * 2, 1, bias=True)
            self.dw = nn.Conv2d(channels * 2, channels * 2, 3, padding=1, groups=channels * 2, bias=True)
            self.sg = SimpleGate()
            self.pw2 = nn.Conv2d(channels, channels, 1, bias=True)
            
            # LayerNorm instead of SCA (matches NAFNet practice)
            # self.norm = LayerNorm2d(channels)
            self.norm = nn.GroupNorm(1, channels)
            
            # FFN: similar structure, but no GELU in main path; GELU only in FFN if needed
            self.ffn1 = nn.Conv2d(channels, channels * 2, 1, bias=True)
            self.ffn_dw = nn.Conv2d(channels * 2, channels * 2, 3, padding=1, groups=channels * 2, bias=True)
            self.ffn_sg = SimpleGate()  # use gate instead of GELU
            self.ffn2 = nn.Conv2d(channels, channels, 1, bias=True)
            
            # Learnable residuals
            self.beta = nn.Parameter(torch.zeros(1, channels, 1, 1))
            self.gamma = nn.Parameter(torch.zeros(1, channels, 1, 1))
    
        def forward(self, x):
            identity = x
    
            # Main branch
            y = self.pw1(x)
            y = self.dw(y)
            y = self.sg(y)
            y = self.pw2(y)
            y = self.norm(y)
            x = identity + self.beta * y
    
            # FFN branch
            identity2 = x
            z = self.ffn1(x)
            z = self.ffn_dw(z)
            z = self.ffn_sg(z)  # gated, no GELU
            z = self.ffn2(z)
            x = identity2 + self.gamma * z
    
            return x
    
    class NAFNet(nn.Module):
        def __init__(self, img_channels=1, width=48, enc_depths=[2,2,4], middle_blocks=8):
            super().__init__()
            self.entry = nn.Conv2d(img_channels, width, 3, padding=1)
            # encoder
            self.encs = nn.ModuleList()
            self.downs = nn.ModuleList()
            ch = width
            for d in enc_depths:
                blocks = nn.Sequential(*[NAFBlockV2(ch) for _ in range(d)])
                self.encs.append(blocks)
                self.downs.append(nn.Conv2d(ch, ch*2, 2, stride=2))  # downsample by strided conv
                ch *= 2
            # middle
            self.middle = nn.Sequential(*[NAFBlockV2(ch) for _ in range(middle_blocks)])
            # decoder
            self.ups = nn.ModuleList()
            self.decs = nn.ModuleList()
            for d in reversed(enc_depths):
                self.ups.append(nn.ConvTranspose2d(ch, ch//2, 2, stride=2))
                ch = ch // 2
                self.decs.append(nn.Sequential(*[NAFBlockV2(ch) for _ in range(d)]))
            # final conv MUST output single channel
            self.exit = nn.Conv2d(width, img_channels, 3, padding=1)
    
        def forward(self, x):
            identity = x
            x = self.entry(x)
            enc_feats = []
            for enc, down in zip(self.encs, self.downs):
                x = enc(x)
                enc_feats.append(x)
                x = down(x)
            x = self.middle(x)
            for up, dec, feat in zip(self.ups, self.decs, reversed(enc_feats)):
                x = up(x)
                # align shapes if needed
                if x.shape[2:] != feat.shape[2:]:
                    # center-crop to min dims
                    min_h = min(x.shape[2], feat.shape[2])
                    min_w = min(x.shape[3], feat.shape[3])
                    x = x[:, :, :min_h, :min_w]
                    feat = feat[:, :, :min_h, :min_w]
                x = x + feat
                x = dec(x)
            out = self.exit(x)
            # return identity - out  # residual subtract -> final shape (B,1,H,W)
            return out
    
    # -------------------------
    # Dataset for self-supervised finetuning
    # -------------------------
    class ImagePatchDataset(Dataset):
        def __init__(self, image, patch_h=128, patch_w=128, num_samples=4000, mask_ratio=0.05):
            self.image = image.astype(np.float32)
            self.H, self.W = image.shape
            self.ph = patch_h
            self.pw = patch_w
            self.num_samples = num_samples
            self.mask_ratio = mask_ratio
            self.coords = [ (np.random.randint(0, max(1, self.H - self.ph + 1)),
                             np.random.randint(0, max(1, self.W - self.pw + 1))) for _ in range(num_samples) ]
        def __len__(self):
            return self.num_samples
        def __getitem__(self, idx):
            y, x = self.coords[idx]
            patch = self.image[y:y+self.ph, x:x+self.pw].copy()
            # create mask
            mask = np.zeros_like(patch, dtype=np.float32)
            num_mask = int(self.mask_ratio * patch.size)
            if num_mask > 0:
                coords = np.random.choice(patch.size, size=num_mask, replace=False)
                flat = mask.ravel()
                flat[coords] = 1.0
                mask = flat.reshape(patch.shape)
            masked_patch = patch.copy()
            masked_patch[mask.astype(bool)] = np.mean(patch)
            # normalize per-patch
            mean = masked_patch.mean()
            std = masked_patch.std() if masked_patch.std() > 1e-6 else 1.0
            inp = (masked_patch - mean) / std
            tgt = (patch - mean) / std
            # inp_t = torch.from_numpy(inp[None]).float()   # (1,H,W)
            inp_t = torch.from_numpy(masked_patch[None]).float()
            # tgt_t = torch.from_numpy(tgt[None]).float()
            tgt_t = torch.from_numpy(patch[None]).float()
            mask_t = torch.from_numpy(mask[None]).float()
            return inp_t, tgt_t, mask_t
    
    # -------------------------
    # Finetune function (self-supervised using masked loss)
    # -------------------------
    def finetune_selfsupervised(model, image, epochs=5, batch_size=16, lr=1e-4, device=DEVICE):
        ds = ImagePatchDataset(image, patch_h=128, patch_w=128, num_samples=3000, mask_ratio=0.05)
        dl = DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY)
        model.train()
        opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-6)
        scaler = torch.amp.GradScaler('cuda',enabled=USE_AMP)
        loss_fn = nn.L1Loss(reduction='none')
        for ep in range(epochs):
            running = 0.0
            cnt = 0
            pbar = tqdm(dl, desc=f"Finetune ep{ep+1}/{epochs}")
            for inp, tgt, mask in pbar:
                inp = inp.to(device); tgt = tgt.to(device); mask = mask.to(device)
                with torch.amp.autocast('cuda', enabled=USE_AMP):
                    out = model(inp)
                    l = loss_fn(out, tgt)
                    masked_loss = (l * mask).sum() / (mask.sum() + 1e-8)
                opt.zero_grad()
                scaler.scale(masked_loss).backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                scaler.step(opt)
                scaler.update()
                running += masked_loss.item()
                cnt += 1
                pbar.set_postfix({"avg_masked_loss": running/cnt})
            print(f"Epoch {ep+1} avg masked loss: {running/cnt:.6f}")
            torch.cuda.empty_cache()
        model.eval()
        return model
    
    # -------------------------
    # Patching & batched inference (memory-safe)
    # -------------------------
    def extract_patches(image, ph, pw, stride_h, stride_w):
        H, W = image.shape
        pad_h = (ph - (H % ph)) % ph
        pad_w = (pw - (W % pw)) % pw
        image_padded = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect') if (pad_h or pad_w) else image.copy()
        Hp, Wp = image_padded.shape
        ys = list(range(0, Hp - ph + 1, stride_h))
        xs = list(range(0, Wp - pw + 1, stride_w))
        if ys[-1] != Hp - ph:
            ys.append(Hp - ph)
        if xs[-1] != Wp - pw:
            xs.append(Wp - pw)
        coords = []
        for y in ys:
            for x in xs:
                coords.append((y, x))
        return image_padded, coords
    
    def batched_patch_inference(image, model, ph, pw, stride_h, stride_w, device, batch_size=2):
        model.eval()
        torch.cuda.empty_cache()
        win = make_hanning_window(ph, pw)
        image_padded, coords = extract_patches(image, ph, pw, stride_h, stride_w)
        Hp, Wp = image_padded.shape
        output = np.zeros_like(image_padded, dtype=np.float32)
        weight = np.zeros_like(image_padded, dtype=np.float32)
    
        num = len(coords)
        idx = 0
        print(f"Processing patches: {num}")
        with torch.no_grad():
            while idx < num:
                # torch.cuda.synchronize()
                bs = min(batch_size, num - idx)
                inp_list = []
                # means = []
                # stds = []
                for j in range(bs):
                    y, x = coords[idx + j]
                    patch = image_padded[y:y+ph, x:x+pw].astype(np.float32)
                    # m = patch.mean(); s = patch.std() if patch.std() > 1e-6 else 1.0
                    # pn = (patch - m) / s
                    pn = patch
                    inp_list.append(pn)
                    # means.append(m); stds.append(s)
                inp_batch = np.stack(inp_list, axis=0)[:, None, :, :]  # (B,1,H,W)
                inp_t = torch.from_numpy(inp_batch).to(device).float()
                if USE_AMP and device.type == 'cuda':
                    torch.cuda.empty_cache()
                    with torch.amp.autocast('cuda'):
                        out = model(inp_t)
                else:
                    out = model(inp_t)
                out_np = out.detach().cpu().numpy()[:, 0, :, :]  # (B,H,W)
                # accumulate
                for j in range(bs):
                    y, x = coords[idx + j]
                    # patch_out = out_np[j] * stds[j] + means[j]
                    patch_out = out_np[j]
                    output[y:y+ph, x:x+pw] += patch_out * win
                    weight[y:y+ph, x:x+pw] += win
                idx += bs
    
        H, W = image.shape
        denoised = output[:H, :W] / (weight[:H, :W] + 1e-8)
        return denoised
    
    def sliding_window_inference(image, model, patch_h, patch_w, stride_h, stride_w, device):
        """
        image: 2D numpy (H, W), float32 in original units
        model: torch model mapping [1,1,H,W] -> [1,1,H,W] (denoised)
        returns: denoised image (H, W)
        """
        model.eval()
        H, W = image.shape
        pad_h = (patch_h - (H % patch_h)) % patch_h
        pad_w = (patch_w - (W % patch_w)) % patch_w
        if pad_h > 0 or pad_w > 0:
            image_padded = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect')
        else:
            image_padded = image
        Hp, Wp = image_padded.shape
    
        # precompute window
        win = make_hanning_window(patch_h, patch_w)
    
        # accumulators
        output = np.zeros_like(image_padded, dtype=np.float32)
        weight = np.zeros_like(image_padded, dtype=np.float32)
    
        ys = list(range(0, Hp - patch_h + 1, stride_h))
        xs = list(range(0, Wp - patch_w + 1, stride_w))
        # ensure last patch covers end
        if ys[-1] != Hp - patch_h:
            ys.append(Hp - patch_h)
        if xs[-1] != Wp - patch_w:
            xs.append(Wp - patch_w)
    
        for y in ys:
            for x in xs:
                patch = image_padded[y:y+patch_h, x:x+patch_w]
                # normalize patch to zero-mean unit-std (helps network generalize)
                mean = patch.mean()
                std = patch.std() if patch.std() > 1e-6 else 1.0
                patch_n = (patch - mean) / std
                inp = torch.from_numpy(patch_n[None, None, :, :]).to(device)
                with torch.no_grad():
                    out = model(inp).cpu().numpy()[0, 0]
                # undo normalization
                out = out * std + mean
                # accumulate with window
                output[y:y+patch_h, x:x+patch_w] += out * win
                weight[y:y+patch_h, x:x+patch_w] += win
    
        # normalize by weight and crop back
        output = output[:H, :W] / (weight[:H, :W] + 1e-8)
        return output
    
    # -------------------------
    # No-GT evaluation helpers
    # -------------------------
    def noise_reduction_rate(I, D):
        return 1 - (np.var(D) / (np.var(I) + 1e-12))
    
    def edge_preservation_index(I, D):
        from scipy import ndimage
        sx = ndimage.sobel(I, axis=1)
        sy = ndimage.sobel(I, axis=0)
        sI = np.hypot(sx, sy)
        sx2 = ndimage.sobel(D, axis=1)
        sy2 = ndimage.sobel(D, axis=0)
        sD = np.hypot(sx2, sy2)
        return sD.sum() / (sI.sum() + 1e-12)
    
    # -------------------------
    # Main runner
    # -------------------------
    def run_pipeline():
        # 1) load image
        # image = load_hdf_as_numpy(UNFILTERED_FILE, UNFILTERED_VAR)
        # np.save('image_unf.npy', image)
        image_raw = np.load('images\\initial_data_unfiltered.npy')
        H, W = image_raw.shape
        print("Loaded image:", image_raw.shape, " dtype:", image_raw.dtype)
        print("Using device:", DEVICE, "AMP enabled:", USE_AMP)

        glob_mean = np.mean(image_raw)
        glob_std = np.std(image_raw)
        if glob_std < 1e-6: glob_std= 1.0
        image = (image_raw - glob_mean) / (glob_std + 1e-8)
    
        # 2) choose model & try loading weights
        if MODEL_NAME.lower() == 'dncnn':
            model = DnCNN(in_channels=1, features=64, depth=12)
            wpath = DN_CNN_WEIGHTS
        elif MODEL_NAME.lower() == 'tinynaf':
            model = TinyNAFNet(in_ch=1, out_ch=1, width=32, num_blocks=8)
            wpath = TINY_NAF_WEIGHTS
        elif MODEL_NAME.lower() == 'nafnet':
            # a moderately-sized NAFNet for 12GB GPU
            model = NAFNet(img_channels=1, width=48, enc_depths=[2,2,4], middle_blocks=8)
            wpath = NAFNET_WEIGHTS
        else:
            raise ValueError("Unknown model")
        model = model.to(DEVICE)
    
        pretrained = False
        if os.path.exists(wpath):
            try:
                model.load_state_dict(torch.load(wpath, map_location=DEVICE))
                print("Loaded pretrained weights from", wpath)
                pretrained = True
            except Exception as e:
                print("Failed loading weights:", e)
    
        # 3) finetune self-supervised if no pretrained
        if (not pretrained) and FINETUNE:
            print("No pretrained weights found — running self-supervised finetuning")
            model = finetune_selfsupervised(model, image, epochs=FINETUNE_EPOCHS, batch_size=FINETUNE_BATCH, lr=LEARNING_RATE, device=DEVICE)
            os.makedirs(WEIGHTS_DIR, exist_ok=True)
            save_path = os.path.join(WEIGHTS_DIR, f"{MODEL_NAME}_finetuned.pth")
            torch.save(model.state_dict(), save_path)
            print("Saved finetuned weights to", save_path)
    
        # 4) batched patch inference on GPU/CPU
        print("Running patch-based batched inference ...")
        torch.backends.cudnn.enabled = False
        torch.backends.cudnn.deterministic = True
        denoised_norm = batched_patch_inference(image, model, PATCH_H, PATCH_W, STRIDE_H, STRIDE_W, DEVICE, batch_size=BATCH_INFERENCE)
        denoised = denoised_norm * glob_std + glob_mean
        # denoised = sliding_window_inference(image, model, PATCH_H, PATCH_W, STRIDE_H, STRIDE_W, DEVICE)

        image = image_raw
        
        # 5) visualize
        diff = image - denoised
        vmin = np.percentile(image, 1)
        vmax = np.percentile(image, 99)
    
        plt.figure(figsize=(18,6))
        plt.subplot(1,3,1)
        plt.imshow(image, cmap='turbo', vmin=vmin, vmax=vmax, aspect='auto')
        plt.title("Original (Unfiltered)")
        plt.axis('off')
    
        plt.subplot(1,3,2)
        plt.imshow(denoised, cmap='turbo', vmin=vmin, vmax=vmax, aspect='auto')
        plt.title(f"Denoised ({MODEL_NAME})")
        plt.axis('off')
    
        plt.subplot(1,3,3)
        dmin = np.percentile(diff, 1); dmax = np.percentile(diff, 99)
        plt.imshow(diff, cmap='seismic', vmin=dmin, vmax=dmax, aspect='auto')
        plt.title("Difference (Original - Denoised)")
        plt.axis('off')
        plt.tight_layout()
        plt.show()
    
        # 6) metrics (no ground truth)
        nrr = noise_reduction_rate(image, denoised)
        epi = edge_preservation_index(image, denoised)
        print(f"Noise reduction rate: {nrr:.4f}")
        print(f"Edge preservation index: {epi:.4f}")
    
        # 7) save output
        outname = f"images\\denoised_{MODEL_NAME}.npy"
        np.save(outname, denoised)
        print("Saved denoised image to", outname)

    run_pipeline()

run_experiment(model_name='dncnn')

Loaded image: (1212, 140)  dtype: float32
Using device: cuda AMP enabled: True
No pretrained weights found — running self-supervised finetuning

Finetune ep1/60: 100%|████████████████████████████████████████| 188/188 [00:10<00:00, 18.67it/s, avg_masked_loss=0.566]

Epoch 1 avg masked loss: 0.566199

Finetune ep2/60: 100%|████████████████████████████████████████| 188/188 [00:09<00:00, 18.91it/s, avg_masked_loss=0.516]

Epoch 2 avg masked loss: 0.515835

Finetune ep3/60: 100%|████████████████████████████████████████| 188/188 [00:09<00:00, 18.89it/s, avg_masked_loss=0.463]

Epoch 3 avg masked loss: 0.462995

Finetune ep4/60: 100%|████████████████████████████████████████| 188/188 [00:10<00:00, 18.78it/s, avg_masked_loss=0.429]

Epoch 4 avg masked loss: 0.428673

Finetune ep5/60: 100%|████████████████████████████████████████| 188/188 [00:10<00:00, 18.68it/s, avg_masked_loss=0.405]

run_experiment(model_name='tinynaf')

Loaded image: (1212, 140)  dtype: float32
Using device: cuda AMP enabled: True
No pretrained weights found — running self-supervised finetuning

Finetune ep1/60: 100%|████████████████████████████████████████| 188/188 [00:23<00:00,  8.00it/s, avg_masked_loss=0.583]

Epoch 1 avg masked loss: 0.583086

Finetune ep2/60: 100%|████████████████████████████████████████| 188/188 [00:23<00:00,  7.90it/s, avg_masked_loss=0.575]

Epoch 2 avg masked loss: 0.575002

Finetune ep3/60: 100%|████████████████████████████████████████| 188/188 [00:23<00:00,  7.97it/s, avg_masked_loss=0.571]

Epoch 3 avg masked loss: 0.571249

Finetune ep4/60: 100%|████████████████████████████████████████| 188/188 [00:22<00:00,  8.35it/s, avg_masked_loss=0.568]

Epoch 4 avg masked loss: 0.567818

Finetune ep5/60: 100%|████████████████████████████████████████| 188/188 [00:22<00:00,  8.35it/s, avg_masked_loss=0.564]

run_experiment(model_name='nafnet')

Loaded image: (1212, 140)  dtype: float32
Using device: cuda AMP enabled: True
No pretrained weights found — running self-supervised finetuning

Finetune ep1/60: 100%|████████████████████████████████████████| 188/188 [01:09<00:00,  2.71it/s, avg_masked_loss=0.576]

Epoch 1 avg masked loss: 0.575674

Finetune ep2/60: 100%|█████████████████████████████████████████| 188/188 [01:09<00:00,  2.69it/s, avg_masked_loss=0.57]

Epoch 2 avg masked loss: 0.569864

Finetune ep3/60: 100%|████████████████████████████████████████| 188/188 [01:09<00:00,  2.72it/s, avg_masked_loss=0.564]

Epoch 3 avg masked loss: 0.564427

Finetune ep4/60: 100%|████████████████████████████████████████| 188/188 [01:07<00:00,  2.79it/s, avg_masked_loss=0.556]

Epoch 4 avg masked loss: 0.555598

Finetune ep5/60: 100%|████████████████████████████████████████| 188/188 [01:09<00:00,  2.71it/s, avg_masked_loss=0.549]

initial_data_unfiltered = np.load('images\\initial_data_unfiltered.npy')

initial_data_filtered = np.load('images\\initial_data_filtered.npy')
baseline_median_data_filtered = np.load('images\\baseline_median_data_filtered.npy')
bm3d_data_filtered = np.load('images\\bm3d_data_filtered.npy')

denoised_dncnn_data_filtered = np.load('images\\denoised_dncnn.npy')
denoised_tinynafnet_data_filtered = np.load('images\\denoised_tinynaf.npy')
denoised_nafnet_data_filtered = np.load('images\\denoised_nafnet.npy')

print(" ############# INITIAL FILTERING ############# ")
print_results(initial_data_filtered, initial_data_unfiltered)

print("\n")
print(" ############# BASELINE MEDIAN FILTERING ############# ")
print_results(baseline_median_data_filtered, initial_data_unfiltered)

print("\n")
print(" ############# BM3D FILTERING ############# ")
print_results(bm3d_data_filtered, initial_data_unfiltered)

print("\n")
print(" ############# DnCNN FILTERING ############# ")
print_results(denoised_dncnn_data_filtered, initial_data_unfiltered)

print("\n")
print(" ############# TinyNAFNet FILTERING ############# ")
print_results(denoised_tinynafnet_data_filtered, initial_data_unfiltered)

print("\n")
print(" ############# NAFNet FILTERING ############# ")
print_results(denoised_nafnet_data_filtered, initial_data_unfiltered)

 ############# INITIAL FILTERING #############


 ############# BASELINE MEDIAN FILTERING #############


 ############# BM3D FILTERING #############


 ############# DnCNN FILTERING #############


 ############# TinyNAFNet FILTERING #############


 ############# NAFNet FILTERING #############

initial_data_unfiltered = np.load('images\\initial_data_unfiltered.npy')

initial_data_filtered = np.load('images\\initial_data_filtered.npy')
baseline_median_data_filtered = np.load('images\\baseline_median_data_filtered.npy')
bm3d_data_filtered = np.load('images\\bm3d_data_filtered.npy')

denoised_dncnn_data_filtered = np.load('images\\denoised_dncnn.npy')
denoised_tinynafnet_data_filtered = np.load('images\\denoised_tinynaf.npy')
denoised_nafnet_data_filtered = np.load('images\\denoised_nafnet.npy')

initial_metrics = evaluate_denoising(initial_data_unfiltered, initial_data_filtered)
baseline_median_metrics = evaluate_denoising(initial_data_unfiltered, baseline_median_data_filtered)
bm3d_metrics = evaluate_denoising(initial_data_unfiltered, bm3d_data_filtered)

dncnn_metrics = evaluate_denoising(initial_data_unfiltered, denoised_dncnn_data_filtered)
tinynafnet_metrics = evaluate_denoising(initial_data_unfiltered, denoised_tinynafnet_data_filtered)
nafnet_metrics = evaluate_denoising(initial_data_unfiltered, denoised_nafnet_data_filtered)

final_metrics = [
    initial_metrics,
    baseline_median_metrics,
    bm3d_metrics,
    dncnn_metrics,
    tinynafnet_metrics,
    nafnet_metrics
]

final_metrics_df = pd.DataFrame(final_metrics)
final_metrics_df.index = [
    'initial_filtering',
    'baseline_median_filter',
    'BM3D',
    'DnCNN',
    'TinyNAFNet',
    'NAFNet'
]

final_metrics_df = final_metrics_df.T
final_metrics_df

	initial_filtering	baseline_median_filter	BM3D	DnCNN	TinyNAFNet	NAFNet
Noise reduction rate (ideal 0.4–0.85)	0.16850	0.17500	0.481000	0.729300	0.577000	0.598000
ENL original	392.21300	392.21300	392.213000	392.213000	392.213000	392.213000
ENL denoised	2589.31300	2028.44900	830.384000	584.428000	508.390000	571.286000
ENL relative increase (ideal > 1.0)	5.60180	4.17180	1.117200	0.490100	0.296200	0.456600
Edge preservation index (ideal 0.8–1.05)	0.09030	0.11430	0.870200	0.733200	0.849200	0.895900
HF energy ratio (ideal 0.3–0.7)	0.00020	0.02360	0.250000	0.708400	0.477100	0.495300
Residual structure SSIM (ideal ≈ 0)	0.45978	0.46736	0.057811	0.835813	0.358322	0.271659

Импорт всех необходимых библиотек¶

Функции для загрузки и визуализации данных¶

Функция, которая считает метрики между исходным и отфильтрованным изображением¶

Посмотрим на исходные данные¶

Baseline: Median Filter¶

Block-matching and 3D filtering (BM3D)¶

DnCNN, NAFNet-mini и NAFNet¶

Визуализация полученных результатов¶

Оценка полученных результатов¶